# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

import os
import datetime
import logging
import sys
import numpy
import pandas

from absl import app, flags

from concurrent.futures import ProcessPoolExecutor
from coin.exchange.okex_futures.kr_rest.futures_product import OkexFuturesProduct
from coin.exchange.huobi_futures.kr_rest.futures_product import HuobiFuturesProduct
from coin.support.proto_log.logic.util import run_from_og_log_archive
from coin.exchange.kr_rest.product.product_impl import generate_product_from_str
from coin.strategy.mm.tool.order_reader import OrderReader, OrderInfo
from coin.strategy.mm.fastfeature.ridge_regress import get_quantile

import coin.strategy.mm.tool.archive_base as abase
import coin.strategy.mm.feed as feed
import coin.strategy.mm.subscription as sub

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

FLAGS = flags.FLAGS


class OrderstatIntraday(object):
  def __init__(self, trading_date, product, nametuple):
    self._nametuple = nametuple
    self._trading_date = trading_date
    self._product = product
    self._order_reader = OrderReader(self._product)

  def analyze(self):
    # moving fill ratio plot
    # order reject ratio pot
    orders_map = self._order_reader.get_orders_map()
    try:
      orderdf = pandas.DataFrame(list(orders_map.values()), columns=OrderInfo.__fields__)
    except Exception:
      orderdf = pandas.DataFrame(list(orders_map.values()), columns=OrderInfo._fields)

    os.makedirs("oe_stat_intra", exist_ok=True)

    baldf = pandas.DataFrame(self._order_reader._balances[self._product.base.currency],
                             columns=['timestamp', 'balance'])
    baldf['currency'] = self._product.base.currency
    posdf = pandas.DataFrame(self._order_reader._balances[self._product.symbol],
                             columns=['timestamp', 'pos'])

    try:
      combdf = pandas.concat(
          [baldf, posdf],
          sort=False).sort_values('timestamp').fillna(method='ffill').reset_index(drop=True)
    except Exception:
      combdf = pandas.concat(
          [baldf, posdf]).sort_values('timestamp').fillna(method='ffill').reset_index(drop=True)

    combdf['ts_sec'] = (combdf['timestamp'] * 1e-9).astype(int)
    combsecdf = combdf.groupby('ts_sec').last().reset_index()
    combsecdf.to_csv(f'oe_stat_intra/{self._nametuple}.csv')

    plt.rcParams['figure.figsize'] = 16, 7
    plt.rcParams['font.family'] = ['Dejavu Sans Mono']
    plt.rcParams['font.sans-serif'] = ['Dejavu Sans Mono']
    plt.rcParams['font.monospace'] = ['Dejavu Sans Mono']
    plt.rcParams['font.size'] = 10
    plt.rcParams['legend.fontsize'] = 10
    plt.rcParams['xtick.labelsize'] = 10
    plt.rcParams['ytick.labelsize'] = 10

    # plt.plot(tidx, ((orderdf['rejected_perhaps_close'] == False) & (orderdf['rejected_timestamp'] > 0)).cumsum(), 'g-', tidx, (orderdf['cancelled_timestamp'] > 0).cumsum(), 'g-', tidx, (orderdf['filled_timestamp'] > 0).cumsum(), 'r-', lw=0.5)
    # plt.legend(['reject msg cumsum', 'cancel msg cumsum', 'fill msg cumsum'], loc='upper left')

    tidx = pandas.DatetimeIndex(orderdf['last_timestamp'])
    plt.plot(tidx, ((orderdf['rejected_perhaps_close'] == False) &
                    (orderdf['rejected_timestamp'] > 0)).cumsum(),
             'g-',
             numpy.nan,
             'r-',
             lw=0.5)
    plt.ylabel("# reject msg")
    plt.legend(['# reject msg', f'{self._product.base.currency} balance'], loc='upper left')
    plt.minorticks_on()
    plt.grid(lw=0.5, which='major')
    plt.grid(lw=0.2, which='minor')
    plt.twinx()
    plt.plot(pandas.DatetimeIndex(baldf['timestamp']), baldf['balance'], 'r-', lw=0.5)
    plt.ylabel(f'{self._product.base.currency} balance')
    plt.title(self._nametuple)
    plt.savefig(f'oe_stat_intra/{self._nametuple}.png')
    plt.close()

  def on_log(self, *args, **kwargs):
    self._order_reader.on_log2(*args, **kwargs)


def launch(baseflags, fsflags, osflags, symbol, postfix):
  product = abase.symbols_to_products([symbol], baseflags)[0]
  start_time, end_time = abase.get_time_range(baseflags)

  nametuple = "_".join((baseflags.exchange,
                        osflags.orderlog_machine,
                        symbol,
                        baseflags.trading_date,
                        osflags.owner or "",
                        postfix or ""))
  strategy = OrderstatIntraday(baseflags.trading_date, product, nametuple)

  abase.run_from_order_archive_base(baseflags, osflags, on_log=strategy.on_log)

  if len(strategy._order_reader._orders_map) == 0:
    logging.info("%s: no order event, skip.", baseflags.trading_date)
    return

  strategy.analyze()
  return 0


def launch_date(trading_date, postfix=None):
  logging.info(trading_date)
  baseflags = abase.get_base_flags()
  baseflags.trading_date = trading_date.strftime("%Y%m%d")
  fsflags = abase.get_feed_subscription_flags()
  osflags = abase.get_order_subscription_flags()
  launch(baseflags, fsflags, osflags, flags.FLAGS.symbol, postfix)


def get_trading_dates(trading_date_str):
  if trading_date_str.find("-") >= 0:
    tdfrom, tdto = trading_date_str.split("-")
    return pandas.date_range(tdfrom, tdto).to_pydatetime()
  else:
    return [datetime.datetime.strptime(trading_date_str, '%Y%m%d')]


def main(_):
  # import warnings; warnings.filterwarnings("error")
  exchange = FLAGS.exchange
  trading_date = FLAGS.trading_date
  machine = FLAGS.orderlog_machine
  assert machine, '--orderlog_machine must be specified.'
  assert exchange, '--exchange must be specified.'
  print('Running for %s %s %s ...' % (trading_date, machine, exchange))
  sys.stdout.flush()
  trading_dates = get_trading_dates(trading_date)
  if len(trading_dates) == 1:
    launch_date(trading_dates[0])
  else:
    with ProcessPoolExecutor(max_workers=FLAGS.cpu) as executor:
      futures = []
      for trading_date in trading_dates:
        futures.append(executor.submit(launch_date, trading_date))
      for future in futures:
        future.result()


if __name__ == '__main__':
  abase.define_base_flags()
  abase.define_order_archive_flags()
  abase.define_feed_archive_flags()

  flags.FLAGS.feed_machine = 'feed-01.ap-northeast-1.aws'
  flags.DEFINE_string('symbol', 'EOS-USD.QUARTER', 'symbol name')

  flags.DEFINE_integer('cpu', 4, 'num cpu to use')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
